# -*- coding: utf-8 -*-
"""
@Author ：mengying
@Date   ：2024/9/3 14:12
@Email  : 652044581@qq.com
@Desc   : 功能描述
"""
import pathlib
import tempfile
import uuid
from pathlib import Path
import cv2
import numpy as np
from PIL import Image
from modelscope.pipelines import pipeline
import os
from modelscope.utils.constant import Tasks
from sanic import Sanic
from sanic.response import file as FileResponse
from sanic.response import json as JsonResponse

app = Sanic("registration-photo")

model = pipeline(Tasks.portrait_matting, model='damo/cv_unet_image-matting')

image_container = pathlib.Path.joinpath(Path(__file__).parent, "image")
if not image_container.exists():
    image_container.mkdir()


def get_temp_path(suffix=".pdf"):
    """获取临时文件路径"""
    _, temp_path = tempfile.mkstemp(suffix=suffix)
    return temp_path


def save_file(fileBody: bytes, suffix=".pdf"):
    """本地保存pdf"""
    temp_path = get_temp_path(suffix=suffix)
    with open(temp_path, 'wb') as f:
        f.write(fileBody)
    return temp_path


def save_image(image_path, img):
    """保存图片"""
    cv2.imencode(".png", img)[1].tofile(image_path)


def hex_to_rgb_255(hex_color):
    """颜色hex 转rgb"""
    if hex_color.startswith('#'):
        hex_color = hex_color[1:]

    if len(hex_color) == 3:
        hex_color = ''.join([c * 2 for c in hex_color])

    r = int(hex_color[0:2], 16)
    g = int(hex_color[2:4], 16)
    b = int(hex_color[4:6], 16)
    return r, g, b


def get_image_path(filename):
    """获取保存的绝对路径"""
    return pathlib.Path.joinpath(image_container, filename).as_posix()


def replace_background_color(image_path, color):
    image = Image.open(image_path)
    if image.mode == 'RGBA':
        image_bg = Image.new('RGB', image.size, color)
        image_bg.paste(image, mask=image.split()[3])
        return image_bg
    else:
        raise Exception('Invalid image')


@app.route('/matting/registration', methods=['POST'])
def matting(request):
    """图片抠图"""
    file = request.files.get("file", "")
    file_path = save_file(file.body, suffix=".png")
    result: np.numarray = model(file_path).get("output_img")
    filename = uuid.uuid4().hex + ".png"
    matting_image = get_image_path(filename)
    save_image(matting_image, result)
    return JsonResponse(body=dict(file_id=filename))


@app.route('/matting/download')
def download(request):
    """图片下载"""
    file_id = request.args.get("file_id", "")
    file_path = get_image_path(file_id)
    return FileResponse(file_path, headers={'Content-Type': 'application/png'})


@app.route('/matting/combine')
def combine(request):
    """图片合成"""
    file_id = request.args.get("file_id", "")
    color = request.args.get("color", "")  # 颜色的hex值
    width = request.args.get("width", "")
    height = request.args.get("height", "")

    if not color:
        return JsonResponse(status=400, body=dict(message="Invalid parameters"))

    r, g, b = hex_to_rgb_255(color)
    file_path = get_image_path(file_id)
    filename = uuid.uuid4().hex + ".png"
    save_path = get_image_path(filename)
    new_image = replace_background_color(file_path, color=(r, g, b))

    # 修改大小
    if width and height:
        resize_image = new_image.resize((int(width), int(height)))
        resize_image.save(save_path)
    else:
        new_image.save(save_path)

    return JsonResponse(body=dict(file_id=filename))


if __name__ == "__main__":
    import platform

    if platform.system().lower() == "windows":
        app.run(host='127.0.0.1', port=8000, auto_reload=True, access_log=True, workers=1)
    else:
        app.run(host='0.0.0.0', port=8000, workers=os.cpu_count())
